import re
import sys
import ctypes
import argparse
import threading
import tkinter as tk
from ctypes import wintypes
from tkinter import ttk, messagebox

# ============================== 常量定义 ==============================
PROCESS_ALL_ACCESS = 0x1F0FFF
THREAD_ALL_ACCESS = 0x1F03FF
THREAD_SUSPEND_RESUME = 0x0002
STATUS_SUCCESS = 0

# ============================== Windows API 初始化 ==============================
kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)
advapi32 = ctypes.WinDLL('advapi32', use_last_error=True)
shell32 = ctypes.WinDLL('shell32', use_last_error=True)

# ============================== 系统结构体定义 ==============================
class PROCESSENTRY32(ctypes.Structure):
    _fields_ = [
        ("dwSize", wintypes.DWORD),
        ("cntUsage", wintypes.DWORD),
        ("th32ProcessID", wintypes.DWORD),
        ("th32DefaultHeapID", ctypes.POINTER(wintypes.ULONG)),
        ("th32ModuleID", wintypes.DWORD),
        ("cntThreads", wintypes.DWORD),
        ("th32ParentProcessID", wintypes.DWORD),
        ("pcPriClassBase", wintypes.LONG),
        ("dwFlags", wintypes.DWORD),
        ("szExeFile", ctypes.c_char * 260)
    ]

class THREADENTRY32(ctypes.Structure):
    _fields_ = [
        ("dwSize", wintypes.DWORD),
        ("cntUsage", wintypes.DWORD),
        ("th32ThreadID", wintypes.DWORD),
        ("th32OwnerProcessID", wintypes.DWORD),
        ("tpBasePri", wintypes.LONG),
        ("tpDeltaPri", wintypes.LONG),
        ("dwFlags", wintypes.DWORD)
    ]

# ============================== Windows API 函数声明 ==============================
kernel32.CreateToolhelp32Snapshot.argtypes = [wintypes.DWORD, wintypes.DWORD]
kernel32.CreateToolhelp32Snapshot.restype = wintypes.HANDLE

kernel32.Process32First.argtypes = [wintypes.HANDLE, ctypes.POINTER(PROCESSENTRY32)]
kernel32.Process32First.restype = wintypes.BOOL
kernel32.Process32Next.argtypes = [wintypes.HANDLE, ctypes.POINTER(PROCESSENTRY32)]
kernel32.Process32Next.restype = wintypes.BOOL

kernel32.Thread32First.argtypes = [wintypes.HANDLE, ctypes.POINTER(THREADENTRY32)]
kernel32.Thread32First.restype = wintypes.BOOL
kernel32.Thread32Next.argtypes = [wintypes.HANDLE, ctypes.POINTER(THREADENTRY32)]
kernel32.Thread32Next.restype = wintypes.BOOL

kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
kernel32.OpenProcess.restype = wintypes.HANDLE
kernel32.TerminateProcess.argtypes = [wintypes.HANDLE, wintypes.UINT]
kernel32.TerminateProcess.restype = wintypes.BOOL

kernel32.OpenThread.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
kernel32.OpenThread.restype = wintypes.HANDLE
kernel32.SuspendThread.argtypes = [wintypes.HANDLE]
kernel32.SuspendThread.restype = wintypes.DWORD
kernel32.ResumeThread.argtypes = [wintypes.HANDLE]
kernel32.ResumeThread.restype = wintypes.DWORD

kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
kernel32.CloseHandle.restype = wintypes.BOOL

shell32.ShellExecuteW.argtypes = [
    wintypes.HWND,
    wintypes.LPCWSTR,
    wintypes.LPCWSTR,
    wintypes.LPCWSTR,
    wintypes.LPCWSTR,
    wintypes.INT
]
shell32.ShellExecuteW.restype = wintypes.HINSTANCE

# ============================== 功能函数实现 ==============================

def get_process_ids_by_name(process_name):
    """通过进程名获取所有匹配的进程ID"""
    process_ids = []
    h_snapshot = kernel32.CreateToolhelp32Snapshot(0x00000002, 0)
    if h_snapshot == -1:
        return process_ids
    
    pe32 = PROCESSENTRY32()
    pe32.dwSize = ctypes.sizeof(PROCESSENTRY32)
    
    if kernel32.Process32First(h_snapshot, ctypes.byref(pe32)):
        while True:
            exe_file = pe32.szExeFile.decode('utf-8', errors='ignore').lower()
            if exe_file == process_name.lower() or exe_file == f"{process_name.lower()}.exe":
                process_ids.append(pe32.th32ProcessID)
            
            if not kernel32.Process32Next(h_snapshot, ctypes.byref(pe32)):
                break
    
    kernel32.CloseHandle(h_snapshot)
    return process_ids

def get_thread_ids(process_id):
    """获取指定进程的所有线程ID"""
    thread_ids = []
    h_snapshot = kernel32.CreateToolhelp32Snapshot(0x00000004, 0)
    if h_snapshot == -1:
        return thread_ids
    
    te32 = THREADENTRY32()
    te32.dwSize = ctypes.sizeof(THREADENTRY32)
    
    if kernel32.Thread32First(h_snapshot, ctypes.byref(te32)):
        while True:
            if te32.th32OwnerProcessID == process_id:
                thread_ids.append(te32.th32ThreadID)
            
            if not kernel32.Thread32Next(h_snapshot, ctypes.byref(te32)):
                break
    
    kernel32.CloseHandle(h_snapshot)
    return thread_ids

def kill_process(process_name):
    """结束指定名称的进程"""
    pids = get_process_ids_by_name(process_name)
    if not pids:
        return False, f"未找到名为 {process_name} 的进程"
    
    for pid in pids:
        h_process = kernel32.OpenProcess(PROCESS_ALL_ACCESS, False, pid)
        if h_process:
            if kernel32.TerminateProcess(h_process, 1):
                kernel32.CloseHandle(h_process)
            else:
                kernel32.CloseHandle(h_process)
                return False, f"无法终止进程: {process_name} (PID: {pid})"
    return True, f"成功终止进程: {process_name}"

def create_process(path, admin=False):
    """创建新进程"""
    try:
        if admin:
            result = shell32.ShellExecuteW(
                None, "runas", path, None, None, 1
            )
            if result <= 32:
                return False, f"无法以管理员权限创建进程. 错误代码: {result}"
            return True, f"以管理员权限启动进程: {path}"
        else:
            startupinfo = wintypes.STARTUPINFOW()
            process_info = wintypes.PROCESS_INFORMATION()
            
            if kernel32.CreateProcessW(
                None,
                path,
                None,
                None,
                False,
                0,
                None,
                None,
                ctypes.byref(startupinfo),
                ctypes.byref(process_info)
            ):
                kernel32.CloseHandle(process_info.hProcess)
                kernel32.CloseHandle(process_info.hThread)
                return True, f"启动进程: {path}"
            return False, "创建进程失败"
    except Exception as e:
        return False, f"创建进程时出错: {str(e)}"

def suspend_process(process_name):
    """挂起指定进程（暂停其所有线程）"""
    pids = get_process_ids_by_name(process_name)
    if not pids:
        return False, f"未找到名为 {process_name} 的进程"
    
    for pid in pids:
        thread_ids = get_thread_ids(pid)
        if not thread_ids:
            return False, f"进程 {process_name} 没有找到线程"
        
        for tid in thread_ids:
            h_thread = kernel32.OpenThread(THREAD_SUSPEND_RESUME, False, tid)
            if h_thread:
                kernel32.SuspendThread(h_thread)
                kernel32.CloseHandle(h_thread)
    return True, f"成功挂起进程: {process_name}"

def resume_process(process_name):
    """恢复被挂起的进程（恢复其所有线程）"""
    pids = get_process_ids_by_name(process_name)
    if not pids:
        return False, f"未找到名为 {process_name} 的进程"
    
    for pid in pids:
        thread_ids = get_thread_ids(pid)
        if not thread_ids:
            return False, f"进程 {process_name} 没有找到线程"
        
        for tid in thread_ids:
            h_thread = kernel32.OpenThread(THREAD_SUSPEND_RESUME, False, tid)
            if h_thread:
                kernel32.ResumeThread(h_thread)
                kernel32.CloseHandle(h_thread)
    return True, f"成功恢复进程: {process_name}"

def get_process_list():
    """获取当前所有进程列表"""
    processes = []
    h_snapshot = kernel32.CreateToolhelp32Snapshot(0x00000002, 0)
    if h_snapshot == -1:
        return processes
    
    pe32 = PROCESSENTRY32()
    pe32.dwSize = ctypes.sizeof(PROCESSENTRY32)
    
    if kernel32.Process32First(h_snapshot, ctypes.byref(pe32)):
        while True:
            exe_file = pe32.szExeFile.decode('utf-8', errors='ignore')
            processes.append({
                'name': exe_file,
                'pid': pe32.th32ProcessID,
                'thread_count': pe32.cntThreads
            })
            if not kernel32.Process32Next(h_snapshot, ctypes.byref(pe32)):
                break
    
    kernel32.CloseHandle(h_snapshot)
    return processes

# ============================== GUI 界面实现 ==============================

class ProcessManagerGUI:
    """进程管理器图形用户界面"""
    
    def __init__(self, root):
        self.root = root
        self.root.title("任务进程管理器")
        self.root.geometry("800x600")
        
        # 主框架
        main_frame = ttk.Frame(root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 控制按钮区域
        control_frame = ttk.LabelFrame(main_frame, text="操作")
        control_frame.pack(fill=tk.X, pady=5)
        
        ttk.Button(control_frame, text="刷新", command=self.refresh_list).pack(side=tk.LEFT, padx=5)
        ttk.Button(control_frame, text="结束进程", command=self.kill_selected).pack(side=tk.LEFT, padx=5)
        ttk.Button(control_frame, text="挂起进程", command=self.suspend_selected).pack(side=tk.LEFT, padx=5)
        ttk.Button(control_frame, text="恢复进程", command=self.resume_selected).pack(side=tk.LEFT, padx=5)
        
        # 进程列表区域
        list_frame = ttk.LabelFrame(main_frame, text="进程列表")
        list_frame.pack(fill=tk.BOTH, expand=True, pady=5)
        
        columns = ("pid", "name", "threads")
        self.tree = ttk.Treeview(
            list_frame, columns=columns, show="headings", selectmode="browse"
        )
        
        self.tree.column("pid", width=100, anchor=tk.CENTER)
        self.tree.column("name", width=400, anchor=tk.W)
        self.tree.column("threads", width=100, anchor=tk.CENTER)
        
        self.tree.heading("pid", text="PID")
        self.tree.heading("name", text="进程名")
        self.tree.heading("threads", text="线程数")
        
        scrollbar = ttk.Scrollbar(list_frame, orient=tk.VERTICAL, command=self.tree.yview)
        self.tree.configure(yscroll=scrollbar.set)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.tree.pack(fill=tk.BOTH, expand=True)
        
        # 启动进程区域
        start_frame = ttk.LabelFrame(main_frame, text="启动新进程")
        start_frame.pack(fill=tk.X, pady=5)
        
        ttk.Label(start_frame, text="路径:").pack(side=tk.LEFT, padx=5)
        self.path_var = tk.StringVar()
        ttk.Entry(start_frame, textvariable=self.path_var, width=50).pack(
            side=tk.LEFT, padx=5, fill=tk.X, expand=True
        )
        
        self.admin_var = tk.BooleanVar()
        ttk.Checkbutton(start_frame, text="管理员权限", variable=self.admin_var).pack(side=tk.LEFT, padx=5)
        
        ttk.Button(start_frame, text="启动", command=self.start_process).pack(side=tk.LEFT, padx=5)
        
        # 状态栏
        self.status_var = tk.StringVar()
        status_bar = ttk.Label(root, textvariable=self.status_var, relief=tk.SUNKEN, anchor=tk.W)
        status_bar.pack(side=tk.BOTTOM, fill=tk.X)
        
        self.refresh_list()
    
    def refresh_list(self):
        """刷新进程列表"""
        self.tree.delete(*self.tree.get_children())
        processes = get_process_list()
        for process in processes:
            self.tree.insert("", tk.END, values=(
                process['pid'],
                process['name'],
                process['thread_count']
            ))
        self.status_var.set(f"已加载 {len(processes)} 个进程")
    
    def get_selected_process(self):
        """获取选中的进程名"""
        selected = self.tree.selection()
        if not selected:
            messagebox.showwarning("警告", "请先选择一个进程")
            return None
        item = self.tree.item(selected[0])
        return item['values'][1]
    
    def kill_selected(self):
        """结束选中的进程"""
        process_name = self.get_selected_process()
        if process_name:
            success, message = kill_process(process_name)
            self.show_result(success, message)
            if success:
                self.refresh_list()
    
    def suspend_selected(self):
        """挂起选中的进程"""
        process_name = self.get_selected_process()
        if process_name:
            success, message = suspend_process(process_name)
            self.show_result(success, message)
    
    def resume_selected(self):
        """恢复选中的进程"""
        process_name = self.get_selected_process()
        if process_name:
            success, message = resume_process(process_name)
            self.show_result(success, message)
    
    def start_process(self):
        """启动新进程"""
        path = self.path_var.get().strip()
        if not path:
            messagebox.showwarning("警告", "请输入可执行文件路径")
            return
        
        admin = self.admin_var.get()
        success, message = create_process(path, admin)
        self.show_result(success, message)
        if success:
            self.path_var.set("")
            self.refresh_list()
    
    def show_result(self, success, message):
        """显示操作结果"""
        if success:
            self.status_var.set(f"操作成功: {message}")
        else:
            messagebox.showerror("错误", message)
            self.status_var.set(f"操作失败: {message}")

# ============================== 命令行解析辅助函数 ==============================

def parse_kv_args(args):
    """
    解析键值对参数
    例如: ['taskname=notepad.exe'] -> {'taskname': 'notepad.exe'}
    """
    result = {}
    for arg in args:
        if '=' in arg:
            key, value = arg.split('=', 1)
            result[key.lower()] = value
    return result

def validate_args(parsed_args, required_keys):
    """验证是否包含必需的参数"""
    for key in required_keys:
        if key not in parsed_args:
            return False, f"缺少必需的参数: {key}"
    return True, ""

# ============================== 主函数 ==============================

def print_detailed_usage():
    """打印详细的使用说明"""
    print("""
任务进程管理器 (Task Process Manager)

功能:
  1. 结束进程        : 终止指定名称的进程
  2. 创建进程        : 启动新的进程（支持管理员权限）
  3. 挂起进程        : 暂停指定进程的所有线程
  4. 恢复进程        : 恢复被挂起的进程
  5. GUI界面        : 启动图形用户界面

使用说明:
  1. 结束进程:
        tpman kill taskname=<进程名>
        示例: tpman kill taskname=notepad.exe

  2. 创建进程:
        tpman start path=<可执行文件路径> admin=<true|false>
        示例: tpman start path="C:\\Windows\\System32\\calc.exe" admin=true

  3. 挂起进程:
        tpman suspend taskname=<进程名>
        示例: tpman suspend taskname=chrome.exe

  4. 恢复进程:
        tpman resume taskname=<进程名>
        示例: tpman resume taskname=chrome.exe

  5. 图形界面:
        tpman gui

注意事项:
  - 操作敏感进程可能需要管理员权限
  - 进程名不区分大小写
  - 支持同时操作多个同名进程
  - 结束系统关键进程可能导致系统不稳定
""")

def main():
    """主函数 - 命令行入口点"""
    if len(sys.argv) == 1:
        print_detailed_usage()
        return
    
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(
        description='任务进程管理器',
        add_help=False
    )
    
    # 添加子命令解析器
    subparsers = parser.add_subparsers(dest='command', help='命令')
    
    # kill 命令
    kill_parser = subparsers.add_parser('kill', help='结束进程')
    kill_parser.add_argument('args', nargs='*', help='参数列表')
    
    # start 命令
    start_parser = subparsers.add_parser('start', help='创建进程')
    start_parser.add_argument('args', nargs='*', help='参数列表')
    
    # suspend 命令
    suspend_parser = subparsers.add_parser('suspend', help='挂起进程')
    suspend_parser.add_argument('args', nargs='*', help='参数列表')
    
    # resume 命令
    resume_parser = subparsers.add_parser('resume', help='恢复进程')
    resume_parser.add_argument('args', nargs='*', help='参数列表')
    
    # gui 命令
    subparsers.add_parser('gui', help='启动图形界面')
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 根据命令执行相应操作
    if args.command == 'kill':
        kv_args = parse_kv_args(args.args)
        valid, msg = validate_args(kv_args, ['taskname'])
        if not valid:
            print(msg)
            return
        success, message = kill_process(kv_args['taskname'])
        print(message)
    
    elif args.command == 'start':
        kv_args = parse_kv_args(args.args)
        valid, msg = validate_args(kv_args, ['path', 'admin'])
        if not valid:
            print(msg)
            return
        
        # 转换admin参数为布尔值
        admin = kv_args['admin'].lower() == 'true'
        success, message = create_process(kv_args['path'], admin)
        print(message)
    
    elif args.command == 'suspend':
        kv_args = parse_kv_args(args.args)
        valid, msg = validate_args(kv_args, ['taskname'])
        if not valid:
            print(msg)
            return
        success, message = suspend_process(kv_args['taskname'])
        print(message)
    
    elif args.command == 'resume':
        kv_args = parse_kv_args(args.args)
        valid, msg = validate_args(kv_args, ['taskname'])
        if not valid:
            print(msg)
            return
        success, message = resume_process(kv_args['taskname'])
        print(message)
    
    elif args.command == 'gui':
        root = tk.Tk()
        app = ProcessManagerGUI(root)
        root.mainloop()

# 程序入口
if __name__ == "__main__":
    main()
